[3/8] Megatron Bridge linear attention QAT and QAD example - #2657

Linear-attention series — 8 PRs

The seven open PRs form native GitHub stack #2806 in this order. #2497 has landed, so #2519 targets main. #2805 contains NVFP4 and is based on #2657; #2788 and #2541 are rebased in sequence. #2503 and #2507 retain their existing commits; their rebases remain deferred until their immediate parents merge.

#2788 owns ReplaySSM and the planned two-precision design; the two-precision design is not implemented yet. #2541 owns vLLM engine integration. Serving-time prefill-GEMM QDQ remains deferred until an optimized fused kernel is available.

What does this PR do?

Type of change: New example.

Add a runnable Megatron Bridge example for recurrent-state QAT and QAD, using #2519's runtime. Bridge owns distributed scheduling, optimization and checkpoints. Both modes use pretrain(config, forward_step); QAD adds a frozen teacher through DistillationProvider.

For prompt/completion data, the example now derives each prefix from collator metadata (context_lengths and token_count, or explicit prefill_lengths and sequence_lengths). It never guesses prompts from loss-mask values. Right padding is excluded; custom packed batches supply per-segment prompt/valid lengths and storage boundaries. Token-only data retains the explicit fixed-split option. Multi-turn data must be rendered into a prompt and selected target continuation before tokenization.

Data-derived phases support TP/DP with PP=1 and CP=1. Fixed-length data retains the existing TP/PP/EP options and CP=1. TP ranks share the phase metadata. The adapter consumes a __next__-only iterator once, and applies the phase to the student only.

The dependency files document Bridge 056b78e / Core e998be07 and pin public vLLM 0.30.0. A KDA model additionally needs a compatible Bridge provider; the real-model workflow evidence is GDN.

Usage

Prompt/completion SFT data: a directory containing training.jsonl records with input and output fields. Without --prefill-tokens, the prompt lengths come from the tokenizer/collator.

bash examples/llm_qat/linear_attention/with_vllm_defaults.sh \
  torchrun --standalone --nproc-per-node=1 examples/llm_qat/linear_attention/train.py \
  --model /path/to/local-model --sft-data /path/to/prompt-data \
  --output /path/to/megatron-qat-checkpoint \
  --recipe general/ptq/linear_attention_state_int8_block32_dynamic \
  --train-steps 10 --length 128

Add --teacher-model /path/to/unquantized-model for QAD. For token-only indexed data, replace --sft-data with --train-data /path/to/train_text_document --prefill-tokens 64. See the example README for export/load instructions and STATE_QUANTIZATION.md for the prefix/decode mismatch and gradient design.

Testing

Rebased onto #2519 at 662fd69622; the parent-relative diff contains the example, tests and documentation.

• DFW H100 job 20217113: QAT and QAD train/resume/export/HF load/Megatron reload tests passed, plus the new real Bridge prompt-derived QAD test. The job also tested the separate NVFP4 slice: four passed in 90.50s. The first attempt failed at single-rank test rendezvous after training; explicit file rendezvous fixed the test, and the complete workflow passed on rerun.
• Focused CPU metadata test covers prompt lengths, right padding, packed boundaries and rejection of invalid/missing lengths. Changed-file pre-commit and diff checks passed.
• Short full-model timing run, DFW 20216405: Qwen3.8-27B, four H100s, TP4/PP1/CP1, sequence128/prefix120/suffix8. Both old and new adapters completed six QAT and six QAD optimizer steps. Four measured-step medians: QAT 1.273→0.794s; QAD 1.201→0.993s. Timings vary and post-update losses differ, so this is throughput/workflow evidence only, not an accuracy or full-model numerical-equivalence claim. This run used the fixed-split entry point.
• Earlier preserved-source GDN runs: ten steps each QAT/QAD on Qwen3.8-27B; save/resume, HF export/load and Megatron reimport passed. Those prior results do not substitute for new-source accuracy evaluation.

No KDA Bridge-provider, long-context quality, or full-engine serving claim is made. The chunk-boundary versus prefix/suffix training-strategy comparison is deferred.

Before your PR is "Ready for review"

• Backward compatibility: additive example; runtime/policy changes belong to [2/8] GDN/KDA state QAT with native forwards and INT8/FP8 state #2519.
• Dependencies/copying: imported Bridge/Core/FLA/vLLM APIs; optional dependency versions documented.
• Necessary tests: shared-fixture QAT/QAD checkpoint round trips, one prompt-derived QAD test and one CPU metadata test.
• Changelog: runtime feature entry belongs to the parent.
• Review: draft; human/Claude review pending.

Additional Information

Step 3/8. State formats and execution are supplied by the parent; NVFP4 and ReplaySSM follow separately. Raw experiment records and the working design walkthrough remain outside the committed example.